
'''
# rag graph流程
'''


import bs4
from langchain import hub
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langgraph.graph import START, StateGraph
from typing_extensions import List, TypedDict

# 2 构建图
# Define prompt for question-answering
prompt = hub.pull("rlm/rag-prompt")

# llm client qwen
import sys
sys.path.append(r'D:\code\other\LLMs\my_langchain')
import keys
from langchain_community.chat_models.tongyi import ChatTongyi
llm = ChatTongyi(streaming=False,) # qwen-turbo

# Define state for application
class State(TypedDict):
    question: str
    context: List[Document]
    answer: str


# Define application steps
from rag_lc import get_rag_query
# node1
rag_query = get_rag_query()
def retrieve(state: State):
    retrieved_docs = rag_query(state["question"]) #  List[Document]
    return {"context": retrieved_docs}

# node2
def generate(state: State):
    docs_content = "\n\n".join(doc.page_content for doc in state["context"])
    messages = prompt.invoke({"question": state["question"], "context": docs_content})
    response = llm.invoke(messages)
    return {"answer": response.content}


# 构建图
graph_builder = StateGraph(State).add_sequence([retrieve, generate])
graph_builder.add_edge(START, "retrieve")
graph = graph_builder.compile()


response = graph.invoke({"question": "点检的代码逻辑是什么"})
response = graph.invoke({"question": "点检的代码逻辑是什么"})
print(response["context"])
print(response["answer"])